import numpy as np
import torch.nn as nn
from transformers import AutoModelForMaskedLM, AutoModel, BertModel
import torch
from torch.optim import Adam, SGD
from scipy.spatial.distance import pdist
from scipy.stats import pearsonr


class SentenceEncoder(nn.Module):
    def __init__(self, config):
        super(SentenceEncoder, self).__init__()
        if 'bert-' in config['pretrain_model_path'].split('/')[-1]:
            # self.encoder = AutoModelForMaskedLM.from_pretrained(config['pretrain_model_path'])
            self.encoder = BertModel.from_pretrained(config['pretrain_model_path'])
        else:
            self.encoder = AutoModel.from_pretrained(config['pretrain_model_path'])
        self.hidden_size = self.encoder.config.hidden_size
        pass
        # self.classify = nn.Linear(hidden_size, config['out_size'])

    def forward(self, x):
        x = self.encoder(x)[0]
        # y = nn.functional.max_pool1d(x.transpose(1, 2), x.shape[1]).squeeze()  # squeeze是把单维度的删除
        # x = nn.functional.max_pool1d(x.transpose(1, 2), x.shape[1]).squeeze(-1)  # squeeze是把单维度的删除
        # x = self.classify(x)
        return x


class Torchmodel(nn.Module):
    def __init__(self, config):
        super(Torchmodel, self).__init__()
        self.sentence_encoder = SentenceEncoder(config)
        self.similarity = nn.CosineSimilarity()
        self.loss = nn.CrossEntropyLoss()
        self.hidden_size = self.sentence_encoder.hidden_size
        self.classify = nn.Linear(self.hidden_size, 5)
        self.pooling_style = config['pooling_style']

    def forward(self, x, label=None):
        x = self.sentence_encoder(x)
        if self.pooling_style == 'max':
            self.pooling_layer = nn.MaxPool1d(x.shape[1])
        else:
            self.pooling_layer = nn.AvgPool1d(x.shape[1])
        x = self.pooling_layer(x.transpose(1, 2)).squeeze()
        predict = self.classify(x)
        if label is not None:  # 维度不同，rel_score为１，怎么算loss
            return self.loss(predict, label.squeeze())
        else:
            return predict


def choose_optimizer(config, model):
    optimizer = config["optimizer"]
    learning_rate = config["learning_rate"]
    if optimizer == "adam":
        return Adam(model.parameters(), lr=learning_rate)
    elif optimizer == "sgd":
        return SGD(model.parameters(), lr=learning_rate)
